import json
import math
from typing import Optional
from typing import List, Dict, Any
from .data_bus import DataBus

from ..mathq import LPF_Fst

class PVT_Ctr:
    def __init__(self, timeStepIn: float, jsonPath: str):
        self.jointNum: int = 0
        self.motor_pos_cur: List[float] = []
        self.motor_pos_des_old: List[float] = []
        self.motor_vel: List[float] = []
        self.motor_tor_out: List[float] = []
        self.motor_pos_des: List[float] = []
        self.motor_vel_des: List[float] = []
        self.motor_tor_des: List[float] = []
        self.pvt_Kp: List[float] = []
        self.pvt_Kd: List[float] = []
        self.maxTor: List[float] = []
        self.maxVel: List[float] = []
        self.maxPos: List[float] = []
        self.minPos: List[float] = []
        self.tau_out_lpf: List[LPF_Fst] = []
        self.PV_enable: List[int] = []
        self.motorName: List[str] = [
            "J_arm_l_01", "J_arm_l_02", "J_arm_l_03", "J_arm_l_04", "J_arm_l_05",
            "J_arm_l_06", "J_arm_l_07", "J_arm_r_01", "J_arm_r_02", "J_arm_r_03",
            "J_arm_r_04", "J_arm_r_05", "J_arm_r_06", "J_arm_r_07",
            "J_head_yaw", "J_head_pitch", "J_waist_pitch", "J_waist_roll", "J_waist_yaw",
            "J_hip_l_roll", "J_hip_l_yaw", "J_hip_l_pitch", "J_knee_l_pitch",
            "J_ankle_l_pitch", "J_ankle_l_roll", "J_hip_r_roll", "J_hip_r_yaw",
            "J_hip_r_pitch", "J_knee_r_pitch", "J_ankle_r_pitch", "J_ankle_r_roll"
        ]

        self.jointNum = len(self.motorName)

        self.tau_out_lpf = [LPF_Fst() for _ in range(self.jointNum)]
        self.motor_vel = [0.0] * self.jointNum
        self.motor_pos_cur = [0.0] * self.jointNum
        self.motor_pos_des_old = [0.0] * self.jointNum
        self.motor_tor_out = [0.0] * self.jointNum
        self.pvt_Kp = [0.0] * self.jointNum
        self.pvt_Kd = [0.0] * self.jointNum
        self.maxTor = [400.0] * self.jointNum
        self.maxVel = [50.0] * self.jointNum
        self.maxPos = [3.14] * self.jointNum
        self.minPos = [-3.14] * self.jointNum
        self.PV_enable = [1] * self.jointNum

        # read joint pvt parameters
        with open(jsonPath, 'r') as f:
            root_read = json.load(f)

        for i in range(self.jointNum):
            self.pvt_Kp[i] = root_read[self.motorName[i]]["kp"]
            self.pvt_Kd[i] = root_read[self.motorName[i]]["kd"]
            self.maxTor[i] = root_read[self.motorName[i]]["maxTorque"]
            self.maxVel[i] = root_read[self.motorName[i]]["maxSpeed"]
            self.maxPos[i] = root_read[self.motorName[i]]["maxPos"]
            self.minPos[i] = root_read[self.motorName[i]]["minPos"]
            fc = root_read[self.motorName[i]]["PVT_LPF_Fc"]
            self.tau_out_lpf[i].setPara(fc, timeStepIn)
            self.tau_out_lpf[i].ftOut(0)

    def dataBusRead(self, busIn: DataBus):
        for i in range(self.jointNum):
            self.motor_pos_cur[i] = busIn.motors_pos_cur[i]
            self.motor_vel[i] = busIn.motors_vel_cur[i]
        self.motor_pos_des = busIn.motors_pos_des
        self.motor_vel_des = busIn.motors_vel_des
        self.motor_tor_des = busIn.motors_tor_des

    def dataBusWrite(self, busIn: DataBus):
        busIn.motors_tor_out = self.motor_tor_out
        busIn.motors_tor_cur = self.motor_tor_out

    def setJointPD(self, kp: float, kd: float, jointName: str):
        try:
            id = self.motorName.index(jointName)
        except ValueError:
            print(f"{jointName} NOT found!")
            return
        self.pvt_Kp[id] = kp
        self.pvt_Kd[id] = kd

    # def calMotorsPVT(self):
    #     for i in range(self.jointNum):
    #         tauDes = 0.0
    #         tauDes = self.PV_enable[i] * self.pvt_Kp[i] * (self.motor_pos_des[i] - self.motor_pos_cur[i]) + \
    #                  self.PV_enable[i] * self.pvt_Kd[i] * (self.motor_vel_des[i] - self.motor_vel[i])
    #         tauDes = self.tau_out_lpf[i].ftOut(tauDes) + self.motor_tor_des[i]
    #         if abs(tauDes) >= abs(self.maxTor[i]):
    #             tauDes = self.sign(tauDes) * self.maxTor[i]
    #         self.motor_tor_out[i] = tauDes
    #         self.motor_pos_des_old[i] = self.motor_pos_des[i]

    def calMotorsPVT(self, deltaP_Lim: Optional[float] = None):
        if deltaP_Lim is None:
            for i in range(self.jointNum):
                tauDes = 0.0
                tauDes = self.PV_enable[i] * self.pvt_Kp[i] * (self.motor_pos_des[i] - self.motor_pos_cur[i]) + \
                        self.PV_enable[i] * self.pvt_Kd[i] * (self.motor_vel_des[i] - self.motor_vel[i])
                tauDes = self.tau_out_lpf[i].ftOut(tauDes) + self.motor_tor_des[i]
                if abs(tauDes) >= abs(self.maxTor[i]):
                    tauDes = self.sign(tauDes) * self.maxTor[i]
                self.motor_tor_out[i] = tauDes
                self.motor_pos_des_old[i] = self.motor_pos_des[i]
        else:    
            for i in range(self.jointNum):
                tauDes = 0.0
                delta = self.motor_pos_des[i] - self.motor_pos_des_old[i]
                if abs(delta) >= abs(deltaP_Lim):
                    delta = deltaP_Lim * self.sign(delta)
                pDes = delta + self.motor_pos_des_old[i]
                tauDes = self.PV_enable[i] * self.pvt_Kp[i] * (pDes - self.motor_pos_cur[i]) + \
                        self.PV_enable[i] * self.pvt_Kd[i] * (self.motor_vel_des[i] - self.motor_vel[i])
                tauDes = self.tau_out_lpf[i].ftOut(tauDes) + self.motor_tor_des[i]
                if abs(tauDes) >= abs(self.maxTor[i]):
                    tauDes = self.sign(tauDes) * self.maxTor[i]
                self.motor_tor_out[i] = tauDes
                self.motor_pos_des_old[i] = pDes

    @staticmethod
    def sign(in_val: float) -> float:
        return 1.0 if in_val >= 0 else -1.0

    def enablePV(self):
        self.PV_enable = [1] * self.jointNum

    def enablePV(self, jtId: int):
        self.PV_enable[jtId] = 1

    def disablePV(self):
        self.PV_enable = [0] * self.jointNum

    def disablePV(self, jtId: int):
        self.PV_enable[jtId] = 0
